from asyncio import sleep
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from uuid import uuid4
from nicegui import ui,events
import os
from pathlib import Path
import cv2
from ultralytics import YOLO
from PIL import Image

class YOLOPredictor:
    def __init__(self):
        self.model = None
        self.source = None
        self.running = False
        self.save_results = False
        self.show_original = True
        self.skip_frames = True
        self.use_gpu = False
        self.native_mode = False
        self.model_files = []
        
    def load_model(self, model_path):
        if model_path:
            self.model = YOLO(model_path)
            ui.notify(f'已加载模型: {Path(model_path).name}')
            
    def get_available_models(self, models_dir='models'):
        """获取模型目录下的所有.pt文件"""
        if not os.path.exists(models_dir):
            os.makedirs(models_dir)
        self.model_files = []
        for file in os.listdir(models_dir):
            if file.endswith(('.pt', '.pth','.onnx')):
                self.model_files.append(os.path.join(models_dir, file))
        print(self.model_files)
    async def start_prediction(self):
        if not self.model or not self.source:
            ui.notify('请先选择模型和输入源')
            return
            
        self.running = True
        if str(self.source).lower().endswith(('.mp4', '.avi', '.mov')):
            await self.process_video()
        else:
            await self.process_image()
            
    async def process_video(self):
        cap = cv2.VideoCapture(str(self.source))
        frame_count = 0
        
        while cap.isOpened() and self.running:
            ret, frame = cap.read()
            if not ret:
                break
                
            if self.skip_frames and frame_count % 25 != 0:
                frame_count += 1
                continue
                    # 创建线程池执行器
            with ThreadPoolExecutor(thread_name_prefix='process_video') as executor:
                # 将预测函数包装成一个偏函数
                predict_func = partial(
                    self.model.predict,
                    source=frame,
                    conf=0.25,
                    device=0 if self.use_gpu else 'cpu'
                )
                
                # 在线程池中执行预测
                future = executor.submit(predict_func)
                # 等待结果但不阻塞事件循环
                while not future.done():
                    await sleep(0.1)
                
                # 获取结果
                results = future.result()    
                # 更新显示
                if self.show_original and frame_count % 25 == 0:
                    self.image_ori.set_source(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))

                if frame_count % 5 == 0:
                    results[0].plot()
                    uuid = str(uuid4())
                    os.makedirs('runs/tmp', exist_ok=True)
                    annotated_image_path = f'runs/tmp/{uuid}.jpg'
                    results[0].save(annotated_image_path)
                    self.image.set_source(annotated_image_path)
                if self.save_results:
                    os.makedirs(f'runs/{os.path.basename(self.source)}', exist_ok=True)
                    results[0].save(f'runs/{os.path.basename(self.source)}/output_{frame_count}.jpg')
                frame_count += 1
            await sleep(0.5)
            
        cap.release()
        self.running = False
    async def process_image(self):
        # 创建线程池执行器
        with ThreadPoolExecutor(thread_name_prefix='process_image') as executor:
            # 将预测函数包装成一个偏函数
            predict_func = partial(
                self.model.predict,
                source=str(self.source),
                conf=0.25,
                device=0 if self.use_gpu else 'cpu'
            )
            
            # 在线程池中执行预测
            future = executor.submit(predict_func)
            # 等待结果但不阻塞事件循环
            while not future.done():
                await sleep(0.1)
            
            # 获取结果
            results = future.result()
            
            # 处理结果
            if self.show_original:
                self.image_ori.set_source(str(self.source))
            
            # 保存和显示结果
            results[0].plot()
            uuid = str(uuid4())
            os.makedirs('runs/tmp', exist_ok=True)
            annotated_image_path = f'runs/tmp/{uuid}.jpg'
            results[0].save(annotated_image_path)
            self.image.set_source(annotated_image_path)
            
            if self.save_results:
                os.makedirs('runs', exist_ok=True)
                results[0].save(f'runs/output_{os.path.basename(self.source)}.jpg')   
            self.running = False 
"""     async def process_image(self):
        #通过子线程处理图片
        
        results = self.model.predict(
            source=str(self.source),
            conf=0.25,
            device=0 if self.use_gpu else 'cpu',
        )
        
        
        if self.show_original:
            self.image_ori.set_source(str(self.source))
        annotated_image = results[0].plot()
        uuid = str(uuid4())
        
        annotated_image=f'runs/tmp/{uuid}.jpg'
        results[0].save(annotated_image)
        self.image.set_source(annotated_image)
        
        if self.save_results:
            results[0].save(f'runs/output_{os.path.basename(self.source)}.jpg')
 """
predictor = YOLOPredictor()
@ui.refreshable
def model_select():
    ui.select(
                options={path: Path(path).name for path in predictor.model_files},
                label='选择模型文件',
                on_change=lambda e: predictor.load_model(e.value)
            ).classes('w-64')
@ui.refreshable
def predictBtn():
    ui.button('开始推理', color=f"{'red' if predictor.running is True else 'blue'}", on_click=lambda: (
                predictor.start_prediction() if not predictor.running 
                else setattr(predictor, 'running', False)
            ))\
            .bind_text_from(predictor,'running',backward=lambda x: '启动' if x is False else '停止')\
            .bind_icon_from(predictor,'running',backward=lambda x: 'play_arrow' if x is False else 'pause').bind_enabled_from(predictor, 'model' ,backward=lambda x: x is not None)
@ui.page('/')
def home():
    with ui.row():
        with ui.card():
            ui.button('加载模型列表', on_click=lambda: {
                predictor.get_available_models(),
                model_select.refresh()
                } )
            ui.label('模型选择').classes('text-h6')
            model_select()

        with ui.card():
            ui.label('输入源选择').classes('text-h6')
            source_upload = ui.upload(
                label='上传图片/视频',
                auto_upload=True,
                on_upload=lambda e: handle_upload(e)
            ).props('accept=.jpg,.jpeg,.png,.mp4,.avi,.mov no-thumbnails' )
        with ui.expansion('选项设置', icon='work'):
            with ui.card():
                ui.switch('保存结果', value=False, on_change=lambda e: setattr(predictor, 'save_results', e.value))
                ui.switch('显示原图', value=True, on_change=lambda e: setattr(predictor, 'show_original', e.value))
                ui.switch('跳帧处理', value=True, on_change=lambda e: setattr(predictor, 'skip_frames', e.value))
                ui.switch('使用GPU', value=False, on_change=lambda e: setattr(predictor, 'use_gpu', e.value))    
        with ui.card():
            predictBtn()
            
        
        def handle_upload(e: events.UploadEventArguments):
            os.makedirs('upload_dir', exist_ok=True)
            os.chdir('upload_dir')
            if os.path.exists(e.name):
                os.remove(e.name)
            with open(e.name, 'wb') as f:
                f.write(e.content.read())
                setattr(predictor, 'source', f'upload_dir/{e.name}')
                os.chdir('../')
    with ui.row().classes('w-full space-x-4'):
        
        # 创建图像显示区域
        predictor.image = ui.image().classes('w-full max-w-[800px] max-h-[600px]')
        predictor.image_ori = ui.image().classes('w-full max-w-[800px] max-h-[600px]')

    

if __name__ == '__main__':
    ui.run( reload=False,show=False, title='YOLO推理工具')